"""
PSO(Particle swarm optimization)
"""

from sko.PSO import PSO
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

plt.style.use(['bmh'])

# Step1：define the problem
def demo_func(x):
    x1, x2, x3 = x
    return x1 ** 2 + (x2 - 0.05) ** 2 + x3 ** 2

# Step2: do PSO
pso = PSO(func=demo_func, n_dim=3, pop=40, max_iter=150, lb=[0, -1, 0.5], ub=[1, 1, 1], w=0.8, c1=0.5, c2=0.5)
pso.run()
print('best_x is ', pso.gbest_x, 'best_y is', pso.gbest_y)

# Step3: plot the result
plt.title('PSO(Particle swarm optimization)')
plt.plot(pso.gbest_y_hist)
plt.savefig('第12章：现代优化算法/PSO.png')
plt.show()
